package com.denlaku.longan.util;

import cn.hutool.core.lang.Pair;
import com.denlaku.longan.core.Current;
import com.denlaku.longan.core.RequestContext;
import com.denlaku.longan.core.RequestContextManager;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.MapUtils;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ForkJoinPool;
import java.util.concurrent.ForkJoinTask;
import java.util.concurrent.ForkJoinWorkerThread;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Supplier;
import java.util.stream.Collectors;

/**
 * @author tianx
 */
@Slf4j
public final class ForkJoinUtil {
    private static class DefaultForkJoinWorkerThread extends ForkJoinWorkerThread {
        protected DefaultForkJoinWorkerThread(ForkJoinPool pool) {
            super(pool);
        }
    }

    private static class DefaultForkJoinWorkerThreadFactory implements ForkJoinPool.ForkJoinWorkerThreadFactory {
        private final AtomicInteger threadNumber = new AtomicInteger();

        @Override
        public ForkJoinWorkerThread newThread(ForkJoinPool pool) {
            DefaultForkJoinWorkerThread thread = new DefaultForkJoinWorkerThread(pool);
            thread.setName("bi-fork-join-poo-" + threadNumber.getAndIncrement());
            return thread;
        }
    }

    private static final int PARALLELISM = Runtime.getRuntime().availableProcessors() << 4;

    private static final DefaultForkJoinWorkerThreadFactory THREAD_FACTORY = new DefaultForkJoinWorkerThreadFactory();

    private static final ForkJoinPool FORK_JOIN_POOL = new ForkJoinPool(PARALLELISM, THREAD_FACTORY, null, false);

    /**
     * 同步执行Action
     *
     * @param action Action
     */
    private static void execute(ForkJoinRunnableAction action) {
        try {
            ForkJoinTask<Void> result = FORK_JOIN_POOL.submit(action);
            result.get();
        } catch (Exception e) {
            log.error("Failed to execute fork join action", e);
        }
    }

    /**
     * 同步执行appliers
     *
     * @param tasks tasks
     * @param throwException 是否抛异常
     */
    public static void execute(List<Runnable> tasks, boolean throwException) {
        RequestContext current = Current.getCurrent();
        Thread runsThread = Thread.currentThread();
        List<Runnable> wrappedTasks = tasks.stream().map(task -> wrapRunnable(task, runsThread, current))
                .collect(Collectors.toList());
        execute(new ForkJoinRunnableAction(wrappedTasks, throwException));
    }

    /**
     * 同步执行appliers
     *
     * @param tasks appliers
     */
    public static void execute(List<Runnable> tasks) {
        execute(tasks, false);
    }

    private static Runnable wrapRunnable(Runnable task, Thread runsThread, RequestContext context) {
        return () -> {
            boolean isRunsThread = Thread.currentThread() == runsThread;
            try {
                if (context != null) {
                    RequestContextManager.setCurrent(context);
                }
                task.run();
            } finally {
                if (isRunsThread) {
                    RequestContextManager.remove();
                }
            }
        };
    }

    /**
     * 同步并发执行Task中所有的子任务，并获取返回值
     *
     * @param task Task
     * @param <V> 返回值类型
     * @return 返回值
     */
    private static <V> List<V> supplyList(ForkJoinSupplierTask<V> task) {
        try {
            ForkJoinTask<List<V>> result = FORK_JOIN_POOL.submit(task);
            return result.get();
        } catch (Exception e) {
            log.error("Failed to submit task", e);
        }
        return Lists.empty();
    }

    /**
     * 同步并发执行suppliers，并获取返回值
     *
     * @param suppliers suppliers
     * @param throwException 是否抛异常
     * @param <V> 返回值类型
     * @return 返回值
     */
    public static <V> List<V> supplyList(List<Supplier<V>> suppliers, boolean throwException) {
        RequestContext current = Current.getCurrent();
        Thread runsThread = Thread.currentThread();
        List<Supplier<V>> wrappedTasks = suppliers.stream().map(task -> wrapSupplier(task, runsThread, current))
                .collect(Collectors.toList());
        return supplyList(new ForkJoinSupplierTask<>(wrappedTasks, throwException));
    }

    /**
     * 同步并发执行suppliers，并获取返回值
     *
     * @param suppliers suppliers
     * @param <V> 返回值类型
     * @return 返回值
     */
    public static <V> List<V> supplyList(List<Supplier<V>> suppliers) {
        return supplyList(suppliers, false);
    }

    private static <V> Supplier<V> wrapSupplier(Supplier<V> task, Thread runsThread, RequestContext context) {
        return () -> {
            try {
                if (context != null) {
                    RequestContextManager.setCurrent(context);
                }
                return task.get();
            } finally {
                if (Thread.currentThread() == runsThread) {
                    RequestContextManager.remove();
                }
            }
        };
    }

    /**
     * 同步并发执行suppliers，并获取返回值
     *
     * @param supplierMap supplierMap
     * @param throwException 是否抛异常
     * @param <K> Key类型
     * @param <V> Value类型
     * @return 返回值
     */
    public static <K, V> Map<K, V> supplyMap(Map<K, Supplier<V>> supplierMap, boolean throwException) {
        if (MapUtils.isEmpty(supplierMap)) {
            return Maps.empty();
        }
        List<Supplier<Pair<K, V>>> supplierWrappers = Lists.ofSize(supplierMap.size());
        supplierMap.forEach((key, supplier) -> {
            Supplier<Pair<K, V>> supplierWrapper = () -> Pair.of(key, supplier.get());
            supplierWrappers.add(supplierWrapper);
        });
        List<Pair<K, V>> list = supplyList(supplierWrappers, throwException);
        Map<K, V> result = Maps.ofSize(supplierMap.size());
        for (Pair<K, V> kvPair : list) {
            result.put(kvPair.getKey(), kvPair.getValue());
        }
        return result;
    }

    /**
     * 同步并发执行suppliers，并获取返回值
     *
     * @param supplierMap supplierMap
     * @param <K> Key类型
     * @param <V> Value类型
     * @return 返回值
     */
    public static <K, V> Map<K, V> supplyMap(Map<K, Supplier<V>> supplierMap) {
        return supplyMap(supplierMap, false);
    }
}
